# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/10 20:07 
@Author  : ZhangShenao 
@File    : vector_store.py 
@Desc    : 向量数据库存储
"""

from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.vectorstores import Qdrant
from langchain_core.vectorstores import VST


class VectorStore:
    """向量数据库存储"""

    def __init__(self):
        pass

    @staticmethod
    def new_store_by_documents(chunked_documents: []) -> VST:
        """
        根据文档创建新的向量数据库存储
        :param chunked_documents: 文档列表
        :return: 向量数据库
        """

        # 创建ZhipuAI Embedding
        embedding = ZhipuAIEmbeddings(model='embedding-2')

        # 创建向量数据库实例
        vectorstore = Qdrant.from_documents(
            documents=chunked_documents,  # 已分块的文档
            embedding=embedding,  # 用OpenAI的Embedding Model做嵌入
            location=':memory:',  # in-memory 模式存储
            collection_name='my_docs', )  # 指定collection_name

        return vectorstore
